# -*- coding: utf-8 -*-
"""
An implementation of the policyValueNet in Tensorflow
Tested in Tensorflow 1.14

@author: Liangyue Jia
"""
import numpy as np
import tensorflow as tf

class PolicyValueNet():
    def __init__(self, dim_num, level_num, model_file=None):
        self.dim_num = dim_num #变量数目
        self.level_num = level_num #变量的水平数
        ## 1. input
        self.input_states = tf.placeholder(tf.float32,
                       shape=[None, 2, dim_num])

        ## 2. Common Networks Layers
        self.commen_flat = tf.reshape(self.input_states, [-1, 2*dim_num])
        self.commen1 = tf.layers.dense(inputs=self.commen_flat,
                                      units=32,
                                      activation=tf.nn.relu)
        self.commen2 = tf.layers.dense(inputs=self.commen1,
                                      units=64,
                                      activation=tf.nn.relu)

        self.commen3 = tf.layers.dense(inputs=self.commen2,
                                      units=128,
                                      activation=tf.nn.relu)
        ## 3. Actor Networks
        self.action_fc1 = tf.layers.dense(inputs=self.commen3,
                                      units=64,
                                      activation=tf.nn.relu)
        self.action_fc2 = tf.layers.dense(inputs=self.action_fc1,
                                      units=level_num,
                                      activation=tf.nn.log_softmax)

        ## 4 Evaluation Networks
        self.evaluation_fc1 = tf.layers.dense(inputs=self.commen3,
                                      units=32,
                                      activation=tf.nn.relu)

        self.evaluation_fc2 = tf.layers.dense(inputs=self.evaluation_fc1,
                                      units=1,
                                      activation=tf.nn.tanh)

        ## 5 Define loss function
        # 5.1 Label1：最终函数值
        # 误差1：均方根误差
        self.labels = tf.placeholder(tf.float32, shape=[None,1])
        self.value_loss = tf.losses.mean_squared_error(self.labels,
                                                       self.evaluation_fc2)

        # 5.2 Label2: 再问题的每一个状态下，执行动作的概率,由MCTS_player提供
        self.mcts_probs = tf.placeholder(
            tf.float32, shape=[None, level_num]) #包含每一个value的选择概率
        #  误差2：category cross entropy
        self.policy_loss = tf.negative(tf.reduce_mean(
            tf.reduce_sum(tf.multiply(self.mcts_probs, self.action_fc2), 1)))  ##让policy与MCTS的预测输出概率一致

        # 5.3 L2 penalty
        l2_penalty_beta = 1e-4
        vars = tf.trainable_variables()
        l2_penalty = l2_penalty_beta * tf.add_n(
            [tf.nn.l2_loss(v) for v in vars if "bias" not in v.name.lower()])

        # 5-4 Add up to be the Loss function
        self.loss = self.value_loss + self.policy_loss + l2_penalty

        ## 6 optimizer
        self.lr = tf.placeholder(tf.float32)
        self.optimizer = tf.train.AdamOptimizer(
            learning_rate=self.lr).minimize(self.loss)
        ## 7 定义session
        self.session = tf.Session()

        # calc policy entropy, for monitoring only
        self.entropy = tf.negative(tf.reduce_mean(
                tf.reduce_sum(tf.exp(self.action_fc2) * self.action_fc2, 1)))

        init = tf.global_variables_initializer()
        self.session.run(init)

        ## 8 保存模型参数
        self.saver = tf.train.Saver()
        if model_file is not None:
            self.restore_mode(model_file)

    def restore_mode(self, model_path):
        self.saver.restore(self.session, model_path)

    def save_model(self, model_path):
        self.saver.save(self.session, model_path)

    def policy_value(self, state_batch):
        """
        :param state_batch: a batch of states
        :return: a batch of action probabilities and state values
        """
        log_act_probs, value = self.session.run(
            [self.action_fc2, self.evaluation_fc2],
            feed_dict={self.input_states: state_batch})
        act_probs = np.exp(log_act_probs)
        return act_probs, value

    def policy_value_fn(self, srq):
        """
        :param srq: 问题对象
        :return: zip:当前状态下可选的（act,prob）列表
                 value: 当前状态下的价值
        """
        depth, legal_value = srq.get_avaliable_move() #legal_value：[levels , ];
        current_state = np.ascontiguousarray(srq.current_state().reshape(
            -1, 2, self.dim_num))
        act_probs, value = self.policy_value(current_state) #act_probs：[1 X levels]; value:[1 X 1]
        act_probs_zip = zip(legal_value, act_probs[0]) #所有可选动作的概率
        return act_probs_zip, value

    def train_step(self, state_batch, mcts_probs, value_batch, lr):
        """训练过程"""
        value_batch = np.reshape(value_batch, (-1, 1))
        loss, entropy, _ = self.session.run(
            [self.loss, self.entropy, self.optimizer],
            feed_dict={self.input_states: state_batch,
                       self.mcts_probs: mcts_probs,
                       self.labels: value_batch,
                       self.lr: lr})
        return loss, entropy

if __name__ == '__main__':
    pv_net = PolicyValueNet(8)